package main

import (
	"encoding/json"
	"fmt"
	"github.com/gin-contrib/cors"
	"github.com/gin-gonic/gin"
	"net/http"
	"time"
)

func main() {

	// 解决跨域
	r := gin.Default()
	config := cors.DefaultConfig()
	config.AllowAllOrigins = true
	config.AllowCredentials = true
	r.Use(cors.New(config))
	_ = r.SetTrustedProxies([]string{MyConfig.Server.Host})

	r.POST("/hello", helloHandler)
	//r.GET("/hello", helloHandler)

	// 首页路由
	r.LoadHTMLGlob("web/*.html") // HTML模板目录
	r.GET("/", func(c *gin.Context) {
		c.HTML(http.StatusOK, "index.html", gin.H{
			"title": "RAG知识库系统",
			"year":  time.Now().Year(),
		})
	})

	var address = MyConfig.Server.Port
	if len(MyConfig.Server.Host) > 0 {
		address = MyConfig.Server.Host + ":" + MyConfig.Server.Port
	}
	_ = r.Run(address)
}

func helloHandler(c *gin.Context) {
	res := new(ResultJson)
	start := time.Now()

	var body struct {
		Search string `form:"search" json:"search" binding:"required"` // 查询内容
		Format string `form:"format" json:"format"`                    // 响应格式
	}

	// 同时支持GET和POST请求
	if c.Request.Method == http.MethodGet {
		if err := c.ShouldBindQuery(&body); err != nil {
			res.Fail(400, "GET参数异常:"+err.Error())
			c.JSON(http.StatusOK, res)
			return
		}
	} else {
		if err := c.ShouldBindJSON(&body); err != nil {
			res.Fail(400, "JSON参数异常:"+err.Error())
			c.JSON(http.StatusOK, res)
			return
		}
	}

	if body.Search == "" {
		res.Fail(400, "search参数不能为空")
		c.JSON(http.StatusOK, res)
		return
	}

	if body.Format == "" {
		body.Format = "json" // 默认值
	}

	// 加载知识库内容
	knowledge, err := loadKnowledgeBase(MyConfig.KnowledgeBaseDir)
	if err != nil {
		res.Fail(500, "加载知识库失败:"+err.Error())
		c.JSON(http.StatusOK, res)
		return
	}

	messages := []Message{
		{
			Role:    "system",
			Content: MyConfig.LLMAgentNotice, // 主系统提示
		},
	}

	// 添加格式指令（放在主提示之后，知识库之前）
	if body.Format == "json" {
		messages = append(messages, Message{
			Role:    "system",
			Content: "请始终使用以下JSON格式回答：\n{\"answer\": \"简明回答\", \"details\": {\"key1\": \"说明1\", \"key2\": \"说明2\"}}",
		})
	}

	// 添加知识库和用户问题
	messages = append(messages,
		Message{
			Role:    "system",
			Content: "知识库内容:\n" + knowledge,
		},
		Message{
			Role:    "user",
			Content: body.Search,
		},
	)

	fmt.Printf("Final messages: %+v\n", messages)

	req := Request{
		Model:    MyConfig.LLMModel, // 大语言模型
		Stream:   false,
		Messages: messages,    // 助手提示+知识库+问题
		Format:   body.Format, // 使用传入的format参数
	}
	fmt.Printf("Ollama请求参数: %+v\n", req)

	resp, err := talkToOllama(MyConfig.DefaultOllamaUrl, req)
	if err != nil {
		res.Fail(500, "请求模型失败:"+err.Error())
		c.JSON(http.StatusOK, res)
		return
	}
	fmt.Printf("Ollama原始响应: %+v\n", resp)

	// 解析响应内容
	var responseContent interface{} = resp.Message.Content
	if body.Format == "json" {
		var jsonObj interface{}
		if err := json.Unmarshal([]byte(resp.Message.Content), &jsonObj); err == nil {
			responseContent = jsonObj
		} else {
			// 解析失败时返回原始内容
			fmt.Printf("JSON解析失败: %v\n原始内容: %s\n", err, resp.Message.Content)
		}
	}

	var data struct {
		Content  interface{} `json:"content"`             // 答案内容（可能是字符串或JSON对象）
		UsedTime int64       `json:"used_time,omitempty"` // 耗时（毫秒）
	}
	data.Content = responseContent
	data.UsedTime = time.Since(start).Milliseconds() // 转换为毫秒

	res.Success(data)
	c.JSON(http.StatusOK, res)
	return
}
